# app/core/database.py

from pymilvus import connections, Collection, utility, FieldSchema, CollectionSchema, DataType
from config.settings import settings
from config.model_config import ModelConfig

def connect_milvus():
    """连接Milvus数据库"""
    try:
        connections.connect("default", host=settings.MILVUS_HOST, port=settings.MILVUS_PORT)
        print(f"Milvus连接成功: {settings.MILVUS_HOST}:{settings.MILVUS_PORT}")
    except Exception as e:
        print(f"Milvus连接失败: {e}")
        raise

def disconnect_milvus():
    """断开Milvus连接"""
    if "default" in connections.list_connections():
        connections.disconnect("default")
        print("Milvus连接已断开")

def init_collections():
    """初始化Milvus集合"""
    connect_milvus()
    
    # 检查并更新文本集合
    if utility.has_collection(settings.MILVUS_COLLECTION):
        print(f"检查文本集合schema: {settings.MILVUS_COLLECTION}")
        collection = Collection(settings.MILVUS_COLLECTION)
        schema = collection.schema
        field_names = [field.name for field in schema.fields]
        
        # 检查是否缺少新字段
        required_fields = ["page_num", "slide_num", "paragraph_num"]
        missing_fields = [field for field in required_fields if field not in field_names]
        
        if missing_fields:
            print(f"文本集合缺少字段: {missing_fields}，需要重新创建集合")
            print("警告：这将删除现有数据！")
            utility.drop_collection(settings.MILVUS_COLLECTION)
            print(f"已删除旧集合: {settings.MILVUS_COLLECTION}")
        else:
            print(f"文本集合schema已是最新版本")
    
    # 创建或重新创建文本集合
    if not utility.has_collection(settings.MILVUS_COLLECTION):
        print(f"创建文本集合: {settings.MILVUS_COLLECTION}")
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=1024),
            FieldSchema(name="file_type", dtype=DataType.VARCHAR, max_length=50),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=ModelConfig.TEXT_VECTOR_DIMENSION),
            FieldSchema(name="page_num", dtype=DataType.INT64),
            FieldSchema(name="slide_num", dtype=DataType.INT64),
            FieldSchema(name="paragraph_num", dtype=DataType.INT64)
        ]
        schema = CollectionSchema(fields, description="RAG文本知识库")
        collection = Collection(settings.MILVUS_COLLECTION, schema)
        index_params = {"index_type": "IVF_FLAT", "metric_type": "IP", "params": {"nlist": 128}}
        collection.create_index(field_name="vector", index_params=index_params)
        print(f"文本集合创建成功")
    
    # 检查并更新图片集合
    if utility.has_collection(settings.MILVUS_COLLECTION_IMAGE):
        print(f"检查图片集合schema: {settings.MILVUS_COLLECTION_IMAGE}")
        collection = Collection(settings.MILVUS_COLLECTION_IMAGE)
        schema = collection.schema
        field_names = [field.name for field in schema.fields]
        
        # 检查是否缺少新字段
        required_fields = ["timestamp", "frame_index"]
        missing_fields = [field for field in required_fields if field not in field_names]
        
        if missing_fields:
            print(f"图片集合缺少字段: {missing_fields}，需要重新创建集合")
            print("警告：这将删除现有数据！")
            utility.drop_collection(settings.MILVUS_COLLECTION_IMAGE)
            print(f"已删除旧集合: {settings.MILVUS_COLLECTION_IMAGE}")
        else:
            print(f"图片集合schema已是最新版本")
    
    # 创建或重新创建图片集合
    if not utility.has_collection(settings.MILVUS_COLLECTION_IMAGE):
        print(f"创建图片集合: {settings.MILVUS_COLLECTION_IMAGE}")
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=1024),
            FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=1024),
            FieldSchema(name="file_type", dtype=DataType.VARCHAR, max_length=50),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=ModelConfig.IMAGE_VECTOR_DIMENSION),
            FieldSchema(name="timestamp", dtype=DataType.FLOAT),
            FieldSchema(name="frame_index", dtype=DataType.INT64)
        ]
        schema = CollectionSchema(fields, description="RAG图片知识库")
        collection = Collection(settings.MILVUS_COLLECTION_IMAGE, schema)
        index_params = {"index_type": "IVF_FLAT", "metric_type": "IP", "params": {"nlist": 128}}
        collection.create_index(field_name="vector", index_params=index_params)
        print(f"图片集合创建成功")

def get_text_collection(collection_name: str = None) -> Collection:
    """获取文本集合"""
    if collection_name is None:
        collection_name = settings.MILVUS_COLLECTION
    
    if not utility.has_collection(collection_name):
        raise Exception(f"文本集合 {collection_name} 不存在")
    
    collection = Collection(collection_name)
    
    # 检查并创建缺失的索引
    try:
        # 尝试加载集合，如果失败说明索引不存在
        collection.load()
    except Exception as e:
        if "index not found" in str(e):
            print(f"检测到集合 {collection_name} 缺少索引，正在创建...")
            index_params = {"index_type": "IVF_FLAT", "metric_type": "IP", "params": {"nlist": 128}}
            collection.create_index(field_name="vector", index_params=index_params)
            print(f"索引创建完成，重新加载集合...")
    collection.load()
        else:
            raise e
    
    return collection

def get_image_collection(collection_name: str = None) -> Collection:
    """获取图片集合"""
    if collection_name is None:
        collection_name = settings.MILVUS_COLLECTION_IMAGE
    
    if not utility.has_collection(collection_name):
        raise Exception(f"图片集合 {collection_name} 不存在")
    
    collection = Collection(collection_name)
    
    # 检查并创建缺失的索引
    try:
        # 尝试加载集合，如果失败说明索引不存在
        collection.load()
    except Exception as e:
        if "index not found" in str(e):
            print(f"检测到集合 {collection_name} 缺少索引，正在创建...")
            index_params = {"index_type": "IVF_FLAT", "metric_type": "IP", "params": {"nlist": 128}}
            collection.create_index(field_name="vector", index_params=index_params)
            print(f"索引创建完成，重新加载集合...")
    collection.load()
        else:
            raise e
    
    return collection

def get_collection_stats():
    """获取集合统计信息"""
    stats = {}
    
    try:
        # 文本集合统计
        if utility.has_collection(settings.MILVUS_COLLECTION):
            text_collection = Collection(settings.MILVUS_COLLECTION)
            stats["text_entities"] = text_collection.num_entities
        
        # 图片集合统计
        if utility.has_collection(settings.MILVUS_COLLECTION_IMAGE):
            image_collection = Collection(settings.MILVUS_COLLECTION_IMAGE)
            stats["image_entities"] = image_collection.num_entities
            
    except Exception as e:
        print(f"获取集合统计信息失败: {e}")
    
    return stats 